import requests

# 导入secrets模块用于生成安全的随机数
import secrets
import base64
import hashlib
from urllib.parse import urlencode, parse_qs, urlparse
import logging

# 配置日志记录级别为INFO
logging.basicConfig(level=logging.INFO)
# 获取当前模块的日志记录器
logger = logging.getLogger(__name__)


class OAuthClient:
    def __init__(self):
        # 资源服务器也就是MCP服务器的地址
        self.mcp_server_url = "http://localhost:8000"
        self.auth_server_metadata = None

    def generate_pkce_challenge(self):
        code_verifier = secrets.token_urlsafe(32)
        code_challenge = (
            base64.urlsafe_b64encode(hashlib.sha256(code_verifier.encode()).digest())
            .decode()
            .rstrip("=")
        )
        return code_verifier, code_challenge

    def start_auth_flow(self, scope="mcp:tools"):
        try:
            if not self.auth_server_metadata:
                print(f"未发现可用的授权服务器")
                return False
            authorization_endpoint = self.auth_server_metadata.get(
                "authorization_endpoint"
            )
            # 生成一个PKCE挑战码
            code_verifier, code_challenge = self.generate_pkce_challenge()
            # 构造授权请求
            auth_params = {
                "response_type": "code",  # 指定响应的类型，需要返回一个授权码
                "client_id": self.client_id,
                "redirect_uri": "http://localhost:8080/callback",  # 授权服务重定向的url地址
                "scope": scope,
                "state": secrets.token_urlsafe(16),  # 在开始授权的请求生成一个state字段
                "code_challenge": code_challenge,  # 挑战码
                "code_challenge_method": "S256",  # 挑战码的算法
                "resource": self.mcp_server_url,  # 资源服务器的地址
            }
            auth_url = f"{authorization_endpoint}?{urlencode(auth_params)}"

            print(" OAuth授权请求参数:")
            print(f"   响应类型: {auth_params['response_type']}")
            print(f"   客户端ID: {auth_params['client_id']}")
            print(f"   重定向URI: {auth_params['redirect_uri']}")
            print(f"   作用域: {auth_params['scope']}")
            print(f"   state: {auth_params['state']}")
            print(f"   代码挑战: {auth_params['code_challenge'][:16]}...")
            print(f"   代码挑战方法: {auth_params['code_challenge_method']}")
            print(f"   资源指示器: {auth_params['resource']}")
            # 打印授权URL
            print(f"授权URL: {auth_url}")
            return {
                "auth_url": auth_url,
                "state": auth_params["state"],
                "code_verifier": code_verifier,
                "auth_params": auth_params,
            }
        except Exception as e:
            print(f"开始授权流程失败:{str(e)}")
            return False

    def register_client(self, client_name="MCP客户端"):
        try:
            if not self.auth_server_metadata:
                print(f"未发现可用的授权服务器")
                return False
            registration_endpoint = self.auth_server_metadata.get(
                "registration_endpoint"
            )

            # 构造客户端注册请求体
            client_metadata = {
                "client_name": client_name,  # 客户端名称
                "redirect_uris": ["http://localhost:8080/callback"],  # 重定向URI
                "grant_types": ["authorization_code"],  # 授权类型
                "response_types": ["code"],  # 响应类型
                "token_endpoint_auth_method": "client_secret_basic",  # 令牌端点认证方法
                "scope": "mcp:tools mcp:resources",  # 作用域
                "application_type": "web",  # 应用类型
            }
            # 打印注册请求相关信息
            print(f"   发送客户端注册请求到: {registration_endpoint}")
            print(f"   客户端名称: {client_name}")
            print(f"   重定向URI: {client_metadata['redirect_uris']}")
            print(f"   授权类型: {client_metadata['grant_types']}")
            print(f"   作用域: {client_metadata['scope']}")
            response = requests.post(registration_endpoint, json=client_metadata)
            if response.status_code == 201:
                client_info = response.json()
                self.client_id = client_info["client_id"]
                self.client_secret = client_info["client_secret"]
                print("客户端注册成功")
                print(f"客户端ID:{self.client_id}")
                print(f"客户端密钥:{self.client_secret[:8]}...")
                return True
            else:
                print(f"客户端注册失败:{response.status_code} {response.text}")
                return False
        except Exception as e:
            print(f"客户端注册失败:{str(e)}")
            return False

    def discover_authorization_server(self):
        try:
            # 发送GET请求获取资源服务器的元数据
            response = requests.get(
                f"{self.mcp_server_url}/.well-known/oauth-protected-resource"
            )
            if response.status_code == 200:
                # 解析JSON响应体
                resource_metadata = response.json()
                print(f"获取资源标识符:{resource_metadata.get("resource")}")
                auth_servers = resource_metadata.get("authorization_servers", [])
                if auth_servers:
                    auth_server = auth_servers[0]
                    self.auth_server_metadata = auth_server
                    issuer = auth_server.get("issuer")
                    registration_endpoint = auth_server.get("registration_endpoint")
                    authorization_endpoint = auth_server.get("authorization_endpoint")
                    token_endpoint = auth_server.get("token_endpoint")
                    scopes_supported = auth_server.get("scopes_supported")
                    code_challenge_methods_supported = auth_server.get(
                        "code_challenge_methods_supported"
                    )
                    print(f"授权服务地址:{issuer}")
                    print(f"注册地址:{registration_endpoint}")
                    print(f"获取授权码的地址:{authorization_endpoint}")
                    print(f"获取令牌的地址:{token_endpoint}")
                    print(f"支持的作用域:{scopes_supported}")
                    print(f"PKCE:{",".join(code_challenge_methods_supported)}")
                    return True
            else:
                print(f"授权服务器未找到")
                return False

        except Exception as e:
            print(f"发现授权服务器失败:{str(e)}")
            return False

    def get_authorization_code(self, auth_info):
        try:
            auth_url = auth_info["auth_url"]
            response = requests.get(auth_url, allow_redirects=False)
            if response.status_code == 302:
                redirect_url = response.headers.get("Location")
                logger.info(f"redirect_url:{redirect_url}")
                if redirect_url:
                    parsed_url = urlparse(redirect_url)
                    query_params = parse_qs(parsed_url.query)
                    authorization_code = query_params.get("code", [None])[0]
                    state = query_params.get("state", [None])[0]
                    if authorization_code and state == auth_info["state"]:
                        print(f"授权码获取成功:{authorization_code}")
                        return authorization_code
                    else:
                        print(f"获取授权码失败")
                        return None
                else:
                    print(f"获取授权码失败")
                    return None
            else:
                print(f"获取授权码失败")
                return None

        except Exception as e:
            print(f"获取授权码失败:{str(e)}")
            return False

    def exchange_authorization_code(self, authorization_code, code_verifier):
        try:
            token_endpoint = self.auth_server_metadata["token_endpoint"]
            token_request = {
                "grant_type": "authorization_code",  # 授权类型
                "code": authorization_code,  # 授权码
                "redirect_uri": "http://localhost:8080/callback",  # 重定向URI
                "client_id": self.client_id,  # 客户端ID
                "code_verifier": code_verifier,  # PKCE挑战码
                "resource": self.mcp_server_url,  # RFC8707资源指示器
            }
            auth = (self.client_id, self.client_secret)
            response = requests.post(token_endpoint, json=token_request, auth=auth)
            if response.status_code == 200:
                token_response = response.json()
                self.access_token = token_response["access_token"]
                print(f"成功获取到了令牌:{self.access_token}")
                return True
            else:
                print(f"获取令牌失败: {response.status_code} {response.text}")
                return False

        except Exception as e:
            print(f"获取令牌失败:{str(e)}")
            return False

    def mcp_server_access(self):
        try:
            if not self.access_token:
                print("请先去获取令牌")
                return
            headers = {
                "Authorization": f"Bearer {self.access_token}",
                "Accept": "application/json",
                "MCP-Protocol-Version": "2025-06-18",
            }
            init_request = {
                "jsonrpc": "2.0",
                "id": "1",
                "method": "initialize",
                "params": {
                    "protocolVersion": "2025-06-18",
                    "capabilities": {},
                    "clientInfo": {"name": "MCP客户端", "version": "1.0.0"},
                },
            }
            response = requests.post(
                f"{self.mcp_server_url}/mcp", json=init_request, headers=headers
            )
            if response.status_code == 200:
                init_result = response.json()
                print("MCP客户端成功的连接上了MCP服务器", init_result)
                tools_request = {
                    "jsonrpc": "2.0",
                    "id": "2",
                    "method": "tools/list",
                    "params": {},
                }
                tools_response = requests.post(
                    f"{self.mcp_server_url}/mcp", json=tools_request, headers=headers
                )
                if tools_response.status_code == 200:
                    tools_result = tools_response.json()
                    tools_list = tools_result["result"]["tools"]
                    print("工具列表获取成功")
                    for tool in tools_list:
                        print(f" - {tool['name']}:{tool['description']}")
                    return True
                else:
                    logger.error(
                        f"工具列表获取失败:{tools_response.status_code} {tools_response.text}"
                    )

        except Exception as e:
            print(f"访问服务器资源失败:{str(e)}")
            return False

    def run(self):
        # 0.访问MCP服务器获取工具列表
        if not self.mcp_server_access():
            print("无法访问MCP服务器")
            return False
        # 1.发现授权服务器
        if not self.discover_authorization_server():
            print("无法发现授权服务器")
            return False
        # 2.在授权服务器上注册此客户端
        if not self.register_client("mcp客户端"):
            print("无法注册客户端")
            return False
        # 3.获取领取授权码的URL地址
        auth_info = self.start_auth_flow()
        if not auth_info:
            print("无法开始授权")
            return False
        # 4.获取授权码
        authorization_code = self.get_authorization_code(auth_info)
        if not authorization_code:
            print("无法获取授权码")
            return False
        # 5.交换访问令牌
        if not self.exchange_authorization_code(
            authorization_code, auth_info["code_verifier"]
        ):
            print("无法获取令牌")
            return False
        # 6.访问MCP服务器获取工具列表
        if not self.mcp_server_access():
            print("无法访问MCP服务器")
            return False


def main():
    try:
        client = OAuthClient()
        client.run()
    except KeyboardInterrupt:
        print("用户取消")
    except Exception as e:
        print(f"运行过程中发生错误:{str(e)}")


if __name__ == "__main__":
    main()
